package com.yzh.openai.domain.openai.service.channel.model.impl;

import com.alibaba.fastjson.JSON;
import com.yzh.openai.domain.openai.model.aggregates.ChatProcessAggregate;
import com.yzh.openai.domain.openai.model.entity.MessageEntity;
import com.yzh.openai.domain.openai.service.channel.model.ModelTypeSelector;
import com.yzh.zhipu.model.ChatRequest;
import com.yzh.zhipu.model.ChatResponse;
import com.yzh.zhipu.model.args.Prompt;
import com.yzh.zhipu.model.vo.EventType;
import com.yzh.zhipu.model.vo.Model;
import com.yzh.zhipu.session.ChatGlmSession;
import lombok.extern.slf4j.Slf4j;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.annotation.Resource;
import java.util.List;
import java.util.stream.Collectors;
@Slf4j
@Service
public class GlmTextGenerator implements ModelTypeSelector {
    @Resource
    private ChatGlmSession glmSession;
    @Override
    public void doMessageResponse(ChatProcessAggregate chatProcess, ResponseBodyEmitter emitter) {
        List<MessageEntity> msg = chatProcess.getMessages();
        //1.请求消息
        List<Prompt> messages = msg.stream()
                .map(entity -> Prompt.builder()
                        .role(entity.getRole())
                        .content(entity.getContent())
                        .build())
                .collect(Collectors.toList());
        //2.封装参数
        ChatRequest chatRequest = new ChatRequest();
        chatRequest.setModel(chatProcess.getModel());
        //这行代码要根据Model版本来设置
        if (chatProcess.getModel().equals(Model.CHATGLM_TURBO)) {
            chatRequest.setPrompt(messages);
        }else {
            chatRequest.setMessages(messages);
        }

        //3.请求应答
        glmSession.completions(chatRequest, new EventSourceListener() {
            @Override
            public void onEvent(@NotNull EventSource eventSource, @Nullable String id, @Nullable String type, @NotNull String data) {
                ChatResponse chatResponse = JSON.parseObject(data, ChatResponse.class);
                String res = chatResponse.getResult(chatResponse);

                /**
                 * 应答完成，关闭流
                 * 这里v4版本判断结束流只能通过res=“”
                 */
                if (EventType.finish.getCode().equals(type) || res.equals("")) {
                    emitter.complete();
                }
                //发送消息
                try {
                    emitter.send(res);
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
            }
        });
    }
}
